package dataset

import (
	"bufio"
	"bytes"
	"chatgpt/helper"
	"chatgpt/helper/code"
	"chatgpt/logq"
	"chatgpt/model/glm/request"
	"chatgpt/plugin/elastic"
	lock2 "chatgpt/services/lock"
	"github.com/slclub/go-tips"
	"github.com/xuri/excelize/v2"
	"io"
	"mime/multipart"
	"strings"
	"sync"
	"time"
)

func init() {
	UploadLock = lock2.New()
}

type Dataset struct{}

var UploadLock sync.Locker

// 清理ES数据
func (this *Dataset) ClearEs() (err error) {
	err = elastic.Client.Clear()
	if err != nil {
		logq.InfoF("ChatGLM.DATASET.ClearEs Error: ", err)
		return err
	}

	return err
}

// 读取txt文件，返回map格式类型
func (this *Dataset) HandleTxt(fileHeader *multipart.FileHeader) (data []*elastic.Wenda, err error) {
	var qaList []*elastic.Wenda

	file, err := fileHeader.Open()
	if err != nil {
		return data, err
	}

	defer func() {
		if err := file.Close(); err != nil {
			logq.InfoF("ChatGLM.DATASET.HandleTxt: close txt error, fileHeader: %v", fileHeader)
		}
	}()

	thebytes, _ := helper.ConvertFileEncoding(file, helper.CONST_CHARSET_UTF_8)
	reader := bufio.NewReader(bytes.NewReader(thebytes))
	warpNum := 0 // 换行数量
	qaNum := 0   // 问答数量

	var itemList [][]string
	for {
		str, readerErr := reader.ReadString('\n')
		if warpNum == 1 {
			warpNum = 0
			if len(itemList) != 0 {
				qaNum = qaNum + 1
			}
			itemList = append(itemList, []string{})
		}

		// 兼容mac和window换行格式
		if str == "\n" || str == "\r\n" {
			warpNum++
			continue
		}

		if len(itemList) == 0 {
			var addItem []string
			itemList = append(itemList, addItem)
		}

		item := itemList[qaNum]
		item = append(item, str)
		if len(item) > 0 {
			itemList[qaNum] = item
		}

		// 读取到结尾，结束循环
		if readerErr == io.EOF {
			break
		}
	}

	resList, _ := this.HandleNullData(itemList)

	for _, val := range resList {
		if len(val) < 1 {
			continue
		}

		var qaItem elastic.Wenda
		qaItem.Prompt = val[0]
		qaItem.Content = strings.Join(val[1:], "")
		qaList = append(qaList, &qaItem)
	}

	return qaList, err
}

// 读取xlsx文件，返回map格式类型
func (this *Dataset) HandleXlsx(fileHeader *multipart.FileHeader) (data []*elastic.Wenda, err error) {
	var qaList []*elastic.Wenda

	fileio, err := fileHeader.Open()
	if err != nil {
		return data, err
	}

	defer func() {
		if err := fileio.Close(); err != nil {
			logq.InfoF("ChatGLM.DATASET.HandleXlsx: close txt error, fileHeader: %v", fileHeader)
		}
	}()

	file, err := excelize.OpenReader(fileio)
	if err != nil {
		logq.InfoF("ChatGLM.DATASET.HandleXlsx: reader xlsx error, fileHeader: %v", fileHeader)
		return data, err
	}

	sheetMap := file.GetSheetMap()
	if len(sheetMap) < 1 {
		return data, code.CODENO(code.DATASET_EXCEL_SHEET_NIL)
	}

	sheetName := sheetMap[1]
	rows, err := file.GetRows(sheetName)
	if err != nil {
		logq.InfoF("ChatGLM.DATASET.HandleXlsx: file.GetRows error, %v", err)
		return qaList, err
	}

	resList, _ := this.HandleNullData(rows)

	for _, row := range resList {
		if len(row) == 0 {
			continue
		}

		var qaItem elastic.Wenda
		qaItem.Prompt = row[0]

		if len(row) > 1 {
			qaItem.Content = row[1]
		} else {
			qaItem.Content = ""
		}

		qaList = append(qaList, &qaItem)
	}
	return qaList, err
}

// 处理数据集垃圾数据
func (this *Dataset) HandleNullData(data [][]string) (list [][]string, err error) {
	if len(data) == 0 {
		return list, err
	}

	var resList [][]string
	for _, val := range data {
		// 处理问题为空的情况
		if len(val) < 1 || val[0] == "" {
			continue
		}

		nullNum := 0
		for _, value := range val {
			if value == "" || value == "\n" {
				nullNum = nullNum + 1
			}
		}

		if nullNum < len(val) {
			resList = append(resList, val)
		}
	}

	return resList, err
}

// 数据添加至ES库,并返回IDList
func (this *Dataset) DataToEs(data []*elastic.Wenda) (ids []string, err error) {
	var idList []string
	for _, val := range data {
		val.Prompt = strings.Trim(val.Prompt, "\t")
		val.Prompt = strings.Trim(val.Prompt, "\r\n")
		val.Prompt = strings.Trim(val.Prompt, "\n")
		val.Content = strings.Trim(val.Content, "\t")
		val.Content = strings.Trim(val.Content, "\r\n")
		val.Content = strings.Trim(val.Content, "\n")
		val.Prompt = helper.TailTrimEmptyBack(val.Prompt)
		val.Content = helper.TailTrimEmptyBack(val.Content)
		id, addErr := elastic.Client.AddDocument(&val)
		if addErr != nil {
			logq.InfoF("ChatGLM.DATASET.qaListToEs: AddDocument error, %v", addErr)
			return ids, addErr
		}

		idList = append(idList, id)
	}
	return idList, err
}

// 获取数据集结果列表
func (this *Dataset) SelectDatasetList(pageInfo request.DatasetFileList) (list []map[string]any, total int64, err error) {
	resList, total, err := elastic.ClientFile.SelectDocumentList(pageInfo)
	if err != nil {
		return list, total, err
	}

	for key, item := range resList {
		if _, ok := item["id_list"]; ok {
			delete(item, "id_list")
			resList[key] = item
		}
	}

	return resList, total, err
}

// 保存数据集结果到ES
func (this *Dataset) SaveDatasetResToEs(params elastic.WendaFile) (err error) {
	cf := elastic.ClientFile

	time.Sleep(200 * time.Millisecond)
	// 先判断数据集是否存储在es，如果存在，删除旧数据集的数据列表，更新idList
	res, err := cf.GetDocumentByKey("md5", params.MD5)

	createdAt := ""
	for _, val := range res {
		if _val, ok := val.(map[string]any); ok {
			createdAt = tips.String(_val["created_at"])

			// 数据集结果存在, 先删除数据集结果数据
			oldID := tips.String(_val["_id"])
			_ = cf.DeleteDocument(oldID)

			if _idList, idOk := _val["id_list"].([]interface{}); idOk {
				//数据集结果存在, 先删除数据集
				for _, value := range _idList {
					id := tips.String(value)
					_ = elastic.Client.DeleteDocument(id)
				}
			}
		}
	}

	params.CreatedAt = createdAt
	addErr := cf.AddFileDocument(&params)
	logq.InfoF("ChatGLM.DATASET.SaveDatasetResToEs: es存在历史数据,删除历史数据, %v", addErr)
	return err
}

// 删除数据列表和数据集结果
func (this *Dataset) DeleteDatasetAndRes(params request.DeleteDatasetAndRes) (err error) {
	cf := elastic.ClientFile

	time.Sleep(200 * time.Millisecond)
	// 先判断数据集是否存储在es，如果存在，删除旧数据集的数据列表，更新idList
	res, err := cf.GetDocumentByKey("md5", params.MD5)

	if err != nil {
		logq.InfoF("ChatGLM.DATASET.DeleteDatasetAndRes: es查询历史数据异常, %v", err)
		return err
	}

	// 历史数据集结果不存在
	if len(res) == 0 {
		logq.InfoF("ChatGLM.DATASET.DeleteDatasetAndRes, Error: es历史数据不存在, %v", code.DATASET_FILE_NIL_ERROR)
		return code.CODENO(code.DATASET_FILE_NIL_ERROR)
	}

	for _, val := range res {
		if _val, ok := val.(map[string]any); ok {
			resFileName := _val["file_full_name"]

			// 文件信息不匹配
			if resFileName != params.FileName {
				logq.InfoF("ChatGLM.DATASET.DeleteDatasetAndRes, Error, %v", code.DATASET_FILE_INVALID)
				return code.CODENO(code.DATASET_FILE_INVALID)
			}

			if _idList, idOk := _val["id_list"].([]interface{}); idOk {
				//数据集结果存在, 先删除数据集
				for _, value := range _idList {
					id := tips.String(value)
					_ = elastic.Client.DeleteDocument(id)
				}
			}

			// 数据集结果存在, 先删除数据集结果数据
			oldID := tips.String(_val["_id"])
			delErr := cf.DeleteDocument(oldID)

			logq.InfoF("ChatGLM.DATASET.DeleteDatasetAndRes: es存在历史数据,删除历史数据结果, %v", delErr)
		}
	}

	return err
}
